from faker import Faker
import random
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from data_mock.utils import FileUtil, MysqlUtils_AI

fake = Faker('zh_CN')

# ========== 手术相关基础数据 ==========
CANCER_OPERATIONS = [
    ("32.2901", "胸腔镜下肺病损切除术", "四级", "Ⅰ类", "微创手术"),
    ("40.5903", "乳腺癌根治术", "四级", "Ⅱ类", "根治手术"),
    ("43.9902", "胃癌根治术", "四级", "Ⅱ类", "根治手术"),
    ("45.7603", "腹腔镜下结肠癌根治术", "四级", "Ⅱ类", "微创手术"),
    ("50.2201", "肝癌切除术", "四级", "Ⅱ类", "切除手术"),
    ("57.7101", "前列腺癌根治术", "四级", "Ⅱ类", "根治手术")
]

ANESTHESIA_TYPES = ["全身麻醉", "硬膜外麻醉", "腰麻", "局部麻醉", "复合麻醉"]
ANESTHESIA_GRADES = ["ASAⅠ", "ASAⅡ", "ASAⅢ", "ASAⅣ", "ASAⅤ"]
INCISION_TYPES = ["Ⅰ类切口", "Ⅱ类切口", "Ⅲ类切口", "Ⅳ类切口"]
HEALING_GRADES = ["甲级愈合", "乙级愈合", "丙级愈合"]
SURGERY_TYPES = ["急诊手术", "择期手术", "限期手术"]

HOSPITAL_INFO = [
    ("H1001", "北京协和医院", "P01", "东院"),
    ("H1002", "上海瑞金医院", "P02", "总院"),
    ("H1003", "广州中山医院", "P03", "肿瘤分院")
]


# ========== 时间范围配置 ==========
def generate_DAY_RANGE(start_day, end_day):
    start = datetime.strptime(start_day, "%Y-%m-%d")
    end = datetime.strptime(end_day, "%Y-%m-%d")
    days = []
    current = start
    while current <= end:
        days.append(current.strftime("%Y-%m-%d"))
        current += relativedelta(days=1)
    return days


DAY_RANGE = generate_DAY_RANGE("2025-08-01", "2025-08-03")


def generate_records_per_day(day: str):
    date_obj = datetime.strptime(day, '%Y-%m-%d')
    return 10 if date_obj.day % 2 == 0 else 5


# ========== 核心生成函数 ==========
def generate_surgery_records():
    sql_statements = []
    used_primary_keys = set()  # 跟踪已使用的主键组合

    # 初始化计数器
    counters = {
        'patient_id': 2001,
        'visit_sn': 2001,
        'record_sn': 10001,
        'sug_record_sn': 80001
    }

    for day in DAY_RANGE:
        RECORDS_PER_day = generate_records_per_day(day)
        for i in range(1, RECORDS_PER_day + 1):
            # 患者基础信息
            patient_id = f"PT{counters['patient_id']}"
            counters['patient_id'] += 1

            visit_sn = f"VIS{counters['visit_sn']}"
            counters['visit_sn'] += 1

            # 医院信息
            hospital_code, hospital_name, branch_code, branch_name = random.choice(HOSPITAL_INFO)

            # 时间信息
            admit_datetime = fake.date_time_between(
                start_date=datetime.strptime(day, "%Y-%m-%d") - timedelta(days=30),
                end_date=datetime.strptime(day, "%Y-%m-%d")
            ).strftime('%Y-%m-%d %H:%M:%S')

            hospital_days = random.randint(3, 30)
            discharge_datetime = (datetime.strptime(admit_datetime, '%Y-%m-%d %H:%M:%S') +
                                  timedelta(days=hospital_days)).strftime('%Y-%m-%d %H:%M:%S')

            # 生成手术记录（每个患者1-3个手术）
            for surgery_seq in range(1, random.randint(1, 4)):
                # 确保主键唯一
                while True:
                    sug_record_sn = f"SUG{counters['sug_record_sn']}"
                    primary_key = f"HIS-{sug_record_sn}"
                    if primary_key not in used_primary_keys:
                        used_primary_keys.add(primary_key)
                        counters['sug_record_sn'] += 1
                        break
                    counters['sug_record_sn'] += 1

                # 手术信息
                icd9_code, operation_name, op_level, incision_type, surgery_type = random.choice(CANCER_OPERATIONS)
                operation_datetime = (datetime.strptime(admit_datetime, '%Y-%m-%d %H:%M:%S') +
                                      timedelta(days=random.randint(1, 10))).strftime('%Y-%m-%d %H:%M:%S')

                data = {
                    'admit_datetime': admit_datetime,
                    'admit_number': str(random.randint(1, 5)),
                    'anesthesia_classification': random.choice(ANESTHESIA_GRADES),
                    'anesthesia_type': random.choice(ANESTHESIA_TYPES),
                    'anesthesiologist': fake.name(),
                    'discharge_datetime': discharge_datetime,
                    'from_table': 'SURGERY_RECORD',
                    'from_yy_record_id': f"SOURCE_{random.randint(10000, 99999)}",
                    'hospital_code': hospital_code,
                    'hospital_name': hospital_name,
                    'icd9_code': icd9_code,
                    'incision_level': random.choice(HEALING_GRADES),
                    'incision_type': incision_type,
                    'medical_record_no': f"MR{random.randint(100000, 999999)}",
                    'operation_datetime': operation_datetime,
                    'operation_duration': str(random.randint(60, 360)),
                    'operation_level': op_level,
                    'operation_name': operation_name,
                    'operation_seq': str(surgery_seq),
                    'patient_id': patient_id,
                    'patient_id_old': f"OLD_{patient_id}",
                    'record_datetime': operation_datetime,
                    'record_sn': f"REC{counters['record_sn']}",
                    'record_source': 'HIS',
                    'record_status': 1,
                    'record_type': '住院病案',
                    'record_update_datetime': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                    'sug_record_sn': sug_record_sn,
                    'surgery_assistant1': fake.name(),
                    'surgery_assistant2': fake.name() if random.random() > 0.3 else None,
                    'surgery_doctor': fake.name(),
                    'surgery_type': surgery_type,
                    'unplanned_more_operation': random.choice(["是", "否"]),
                    'visit_sn': visit_sn,
                    'yy_collection_datetime': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                    'yy_record_md5': f"md5_{random.getrandbits(128):032x}",
                    'yy_upload_status': 0,
                    'yy_etl_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                    'yy_upload_time': None,
                    'yy_batch_time': day,
                    'yy_record_batch_id': f"BATCH{day}_{i}_{surgery_seq}",
                    'yy_backfill_time': None,
                    'yy_backfill_status': None,
                    'branch_code': branch_code,
                    'branch_name': branch_name,
                    'date_for_partition': operation_datetime
                }

                sql = _generate_sql('b04_1_2', data)
                sql_statements.append(sql)

    return sql_statements


def _generate_sql(table, data):
    columns = ', '.join([f'`{k}`' for k in data.keys()])
    values = []
    for v in data.values():
        if v is None:
            values.append('NULL')
        elif isinstance(v, (int, float)):
            values.append(str(v))
        else:
            escaped_value = str(v).replace("'", "''")
            values.append(f"'{escaped_value}'")
    return f"INSERT INTO `{table}` ({columns}) VALUES ({', '.join(values)});"


# ========== 执行生成 ==========
if __name__ == "__main__":
    records = generate_surgery_records()
    # 写入数据库
    MysqlUtils_AI.insert_data_to_hub(records, 'b04_1_2')
    # 或写入 SQL 文件
    # FileUtil.generate_sql_file(records, "b04_1_2_surgery_records.sql")